import sys
import cv2
import numpy as np
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,QLabel, QPushButton, QSlider, QFileDialog, QSpinBox, QGroupBox)
from PyQt5.QtCore import Qt
from PyQt5.QtGui import QImage, QPixmap

class HSVMaskGenerator(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("HSV图像二值掩码调试工具")
        self.setGeometry(100, 100, 1000, 600)
        
        # 初始化变量
        self.image_path = None
        self.original_image = None
        self.hsv_image = None
        self.mask = None
        
        # 初始化UI
        self.init_ui()
        
        # 默认HSV范围
        self.h_min = 0
        self.h_max = 179
        self.s_min = 0
        self.s_max = 255
        self.v_min = 0
        self.v_max = 255
        
    def init_ui(self):
        # 主窗口部件
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        
        # 修改主布局为垂直布局
        main_layout = QVBoxLayout(central_widget)
        
        # 添加一个水平布局用于图像和控制面板
        content_layout = QHBoxLayout()
        main_layout.addLayout(content_layout)
        
        # 左侧控制面板 (原有代码保持不变)
        control_panel = QGroupBox("控制面板")
        control_layout = QVBoxLayout()
        
        # 图片选择按钮
        self.btn_open = QPushButton("选择图片")
        self.btn_open.clicked.connect(self.open_image)
        control_layout.addWidget(self.btn_open)
        
        # HSV阈值控制
        self.create_hsv_controls(control_layout)
        
        # 保存按钮
        self.btn_save = QPushButton("保存掩码")
        self.btn_save.clicked.connect(self.save_mask)
        self.btn_save.setEnabled(False)
        control_layout.addWidget(self.btn_save)
        
        control_panel.setLayout(control_layout)
        content_layout.addWidget(control_panel)
        
        # 右侧图像显示区域 (原有代码保持不变)
        image_panel = QWidget()
        image_layout = QVBoxLayout()
        
        # 原始图像显示
        self.lbl_original = QLabel("原始图像")
        self.lbl_original.setAlignment(Qt.AlignCenter)
        self.lbl_original.setMinimumSize(640, 360)
        image_layout.addWidget(self.lbl_original)
        
        # 掩码图像显示
        self.lbl_mask = QLabel("二值掩码")
        self.lbl_mask.setAlignment(Qt.AlignCenter)
        self.lbl_mask.setMinimumSize(640, 360)
        image_layout.addWidget(self.lbl_mask)
        
        image_panel.setLayout(image_layout)
        content_layout.addWidget(image_panel)
            
        # 在main_layout的最后添加信息栏 (放在content_layout之后)
        # 分隔线
        main_layout.addWidget(QLabel("<hr>"))
        
        # 软件信息
        developer_info = QLabel(
            "HSV Mask Tool v1.0 | "
            "Developer: +4 | "
            "Copyright © 2025 | "
            "License: MIT License "
        )
        developer_info.setAlignment(Qt.AlignCenter)
        developer_info.setStyleSheet("color: gray; font-size: 10px;")
        main_layout.addWidget(developer_info)

    
    def create_hsv_controls(self, layout):
        # H通道控制
        h_group = QGroupBox("H通道 (0-179)")
        h_layout = QVBoxLayout()
        
        # H最小值
        h_min_layout = QHBoxLayout()
        lbl_h_min = QLabel("最小值:")
        self.slider_h_min = QSlider(Qt.Horizontal)
        self.slider_h_min.setRange(0, 179)
        self.slider_h_min.valueChanged.connect(self.update_hsv)
        self.spin_h_min = QSpinBox()
        self.spin_h_min.setRange(0, 179)
        self.spin_h_min.valueChanged.connect(self.sync_h_min)
        
        h_min_layout.addWidget(lbl_h_min)
        h_min_layout.addWidget(self.slider_h_min)
        h_min_layout.addWidget(self.spin_h_min)
        h_layout.addLayout(h_min_layout)
        
        # H最大值
        h_max_layout = QHBoxLayout()
        lbl_h_max = QLabel("最大值:")
        self.slider_h_max = QSlider(Qt.Horizontal)
        self.slider_h_max.setRange(0, 179)
        self.slider_h_max.setValue(179)
        self.slider_h_max.valueChanged.connect(self.update_hsv)
        self.spin_h_max = QSpinBox()
        self.spin_h_max.setRange(0, 179)
        self.spin_h_max.setValue(179)
        self.spin_h_max.valueChanged.connect(self.sync_h_max)
        
        h_max_layout.addWidget(lbl_h_max)
        h_max_layout.addWidget(self.slider_h_max)
        h_max_layout.addWidget(self.spin_h_max)
        h_layout.addLayout(h_max_layout)
        
        h_group.setLayout(h_layout)
        layout.addWidget(h_group)
        
        # S通道控制
        s_group = QGroupBox("S通道 (0-255)")
        s_layout = QVBoxLayout()
        
        # S最小值
        s_min_layout = QHBoxLayout()
        lbl_s_min = QLabel("最小值:")
        self.slider_s_min = QSlider(Qt.Horizontal)
        self.slider_s_min.setRange(0, 255)
        self.slider_s_min.valueChanged.connect(self.update_hsv)
        self.spin_s_min = QSpinBox()
        self.spin_s_min.setRange(0, 255)
        self.spin_s_min.valueChanged.connect(self.sync_s_min)
        
        s_min_layout.addWidget(lbl_s_min)
        s_min_layout.addWidget(self.slider_s_min)
        s_min_layout.addWidget(self.spin_s_min)
        s_layout.addLayout(s_min_layout)
        
        # S最大值
        s_max_layout = QHBoxLayout()
        lbl_s_max = QLabel("最大值:")
        self.slider_s_max = QSlider(Qt.Horizontal)
        self.slider_s_max.setRange(0, 255)
        self.slider_s_max.setValue(255)
        self.slider_s_max.valueChanged.connect(self.update_hsv)
        self.spin_s_max = QSpinBox()
        self.spin_s_max.setRange(0, 255)
        self.spin_s_max.setValue(255)
        self.spin_s_max.valueChanged.connect(self.sync_s_max)
        
        s_max_layout.addWidget(lbl_s_max)
        s_max_layout.addWidget(self.slider_s_max)
        s_max_layout.addWidget(self.spin_s_max)
        s_layout.addLayout(s_max_layout)
        
        s_group.setLayout(s_layout)
        layout.addWidget(s_group)
        
        # V通道控制
        v_group = QGroupBox("V通道 (0-255)")
        v_layout = QVBoxLayout()
        
        # V最小值
        v_min_layout = QHBoxLayout()
        lbl_v_min = QLabel("最小值:")
        self.slider_v_min = QSlider(Qt.Horizontal)
        self.slider_v_min.setRange(0, 255)
        self.slider_v_min.valueChanged.connect(self.update_hsv)
        self.spin_v_min = QSpinBox()
        self.spin_v_min.setRange(0, 255)
        self.spin_v_min.valueChanged.connect(self.sync_v_min)
        
        v_min_layout.addWidget(lbl_v_min)
        v_min_layout.addWidget(self.slider_v_min)
        v_min_layout.addWidget(self.spin_v_min)
        v_layout.addLayout(v_min_layout)
        
        # V最大值
        v_max_layout = QHBoxLayout()
        lbl_v_max = QLabel("最大值:")
        self.slider_v_max = QSlider(Qt.Horizontal)
        self.slider_v_max.setRange(0, 255)
        self.slider_v_max.setValue(255)
        self.slider_v_max.valueChanged.connect(self.update_hsv)
        self.spin_v_max = QSpinBox()
        self.spin_v_max.setRange(0, 255)
        self.spin_v_max.setValue(255)
        self.spin_v_max.valueChanged.connect(self.sync_v_max)
        
        v_max_layout.addWidget(lbl_v_max)
        v_max_layout.addWidget(self.slider_v_max)
        v_max_layout.addWidget(self.spin_v_max)
        v_layout.addLayout(v_max_layout)
        
        v_group.setLayout(v_layout)
        layout.addWidget(v_group)
    
    def sync_h_min(self, value):
        self.slider_h_min.setValue(value)
        self.h_min = value
        self.update_hsv()
    
    def sync_h_max(self, value):
        self.slider_h_max.setValue(value)
        self.h_max = value
        self.update_hsv()
    
    def sync_s_min(self, value):
        self.slider_s_min.setValue(value)
        self.s_min = value
        self.update_hsv()
    
    def sync_s_max(self, value):
        self.slider_s_max.setValue(value)
        self.s_max = value
        self.update_hsv()
    
    def sync_v_min(self, value):
        self.slider_v_min.setValue(value)
        self.v_min = value
        self.update_hsv()
    
    def sync_v_max(self, value):
        self.slider_v_max.setValue(value)
        self.v_max = value
        self.update_hsv()
    
    def open_image(self):
        options = QFileDialog.Options()
        file_path, _ = QFileDialog.getOpenFileName(
            self, "选择图片", "", 
            "图片文件 (*.jpg *.jpeg *.png *.bmp);;所有文件 (*)", 
            options=options
        )
        
        if file_path:
            self.image_path = file_path
            self.original_image = cv2.imread(file_path)
            if self.original_image is not None:
                # 转换为HSV颜色空间
                self.hsv_image = cv2.cvtColor(self.original_image, cv2.COLOR_BGR2HSV)
                # 显示原始图像
                self.display_image(self.original_image, self.lbl_original)
                # 启用保存按钮
                self.btn_save.setEnabled(True)
                # 初始掩码
                self.update_hsv()
    
    def update_hsv(self):
        if self.hsv_image is not None:
            # 获取当前HSV阈值
            h_min = self.slider_h_min.value()
            h_max = self.slider_h_max.value()
            s_min = self.slider_s_min.value()
            s_max = self.slider_s_max.value()
            v_min = self.slider_v_min.value()
            v_max = self.slider_v_max.value()
            
            # 更新spin box值
            self.spin_h_min.setValue(h_min)
            self.spin_h_max.setValue(h_max)
            self.spin_s_min.setValue(s_min)
            self.spin_s_max.setValue(s_max)
            self.spin_v_min.setValue(v_min)
            self.spin_v_max.setValue(v_max)
            
            # 创建HSV范围
            lower = np.array([h_min, s_min, v_min])
            upper = np.array([h_max, s_max, v_max])
            
            # 生成掩码
            self.mask = cv2.inRange(self.hsv_image, lower, upper)
            
            # 显示掩码
            self.display_image(self.mask, self.lbl_mask)
    
    def display_image(self, image, label):
        if len(image.shape) == 2:  # 灰度或二值图像
            h, w = image.shape
            bytes_per_line = w
            q_img = QImage(image.data, w, h, bytes_per_line, QImage.Format_Grayscale8)
        else:  # 彩色图像
            h, w, ch = image.shape
            bytes_per_line = ch * w
            # OpenCV使用BGR格式，需要转换为RGB
            rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            q_img = QImage(rgb_image.data, w, h, bytes_per_line, QImage.Format_RGB888)
        
        label.setPixmap(QPixmap.fromImage(q_img).scaled(
            label.width(), label.height(), Qt.KeepAspectRatio
        ))
    
    def save_mask(self):
        if self.mask is not None:
            options = QFileDialog.Options()
            file_path, _ = QFileDialog.getSaveFileName(
                self, "保存掩码", "", 
                "PNG图像 (*.png);;JPEG图像 (*.jpg *.jpeg);;BMP图像 (*.bmp);;所有文件 (*)", 
                options=options
            )
            
            if file_path:
                cv2.imwrite(file_path, self.mask)

if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = HSVMaskGenerator()
    window.show()
    sys.exit(app.exec_())